import numpy as np
import pandas as pd
from matplotlib import pyplot as  plt
from sklearn.model_selection import train_test_split
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import Dataset
from torch.utils.data import DataLoader

"""
1. ✨✨✨ 准备数据✨✨✨ 
   pytorch要求对数据进行打包，打包数据的方法分两步：
    a.重载Dataset类。
    b.DataLoader 数据加载器。
"""
#重载dataset类
class MyDataset(Dataset):
   def __init__(self,X,y):
       self.X = X
       self.y = y
   #重载__getitem__方法。
   def __getitem__(self, idx):
       X = self.X[idx]
       y = self.y[idx]
       return torch.tensor(data=self.X[idx]).float(),torch.tensor(data=self.y[idx]).float()
   def __len__(self):
       return len(self.X)

"""
2.✨✨✨ 定义模型✨✨✨ 
"""
class LinearRegression(nn.Module):
   def __init__(self,in_featurs = 13,out_features = 1):
       super(LinearRegression,self).__init__() #一个子类继承一个父类，所以子类在构建的时候，要先构建一个父类。此行代码写法固定。
       self.linear1 = nn.Linear(in_features=in_featurs,out_features=5)
       self.linear2 = nn.Linear(in_features=5,out_features=out_features)

   def forward(self,x):
       """
       定义正向传播的过程。
       当参数已经训练好了以后，正向传播就是正向推理即predict。
       x的形状必须是[b,in_featurs],out的形状是[b,1]。b是batchsize。
       """
    #    print(x)
       h1 = self.linear1(x)  # 批量的计算。
       h2 = F.relu(h1)
       out = self.linear2(h2)  # 批量的计算。
       return out

# 调试代码。参数定义和业务逻辑都定义在层里面了。我们现在可以直接来看看模型参数：
#print(lr.parameters())
#for ele in lr.parameters():
#   print(ele)

"""
3.✨✨✨ 模型训练✨✨✨ 
模型训练包括：
指定损失函数；
指定优化器；
进行梯度下降
"""
class Trainer:
    def __init__(self, model, train_dataloader, test_dataloader, loss_fn, optimizer, device='cpu'):
        """
        初始化Trainer类

        参数:
        model (nn.Module): PyTorch模型
        train_dataloader (DataLoader): 训练集的数据加载器
        test_dataloader (DataLoader): 测试集的数据加载器
        loss_fn (nn.Module): 损失函数
        optimizer (torch.optim.Optimizer): 优化器
        device (str): 设备, 默认为'cpu'，可以设置为'cuda'来使用GPU
        """
        self.model = model.to(device)  # 模型放置到指定设备
        self.train_dataloader = train_dataloader
        self.test_dataloader = test_dataloader
        self.loss_fn = loss_fn
        self.optimizer = optimizer
        self.device = device

    def get_loss(self, dataloader):
        """
        计算给定数据集上的平均损失

        参数:
        dataloader (DataLoader): 数据加载器

        返回:
        float: 数据集的平均损失
        """
        self.model.eval()  # 评估模式，不进行梯度计算
        losses = []
        with torch.no_grad():  # 禁用梯度计算
            for X, y in dataloader:
                X = X.to(self.device)
                y = y.to(self.device)
                y_pred = self.model(X)  # 前向传播
                loss = self.loss_fn(y_pred, y)  # 计算损失
                losses.append(loss.item())  # 获取损失值
        return np.mean(losses)  # 返回平均损失

    def train(self, epochs=20):
        """
        开始训练模型

        参数:
        epochs (int): 训练轮数，默认为20
        """
        for epoch in range(1, epochs + 1):
            self.model.train()  # 模型进入训练模式
            for X, y in self.train_dataloader:
                X = X.to(self.device)
                y = y.to(self.device)
                
                # 前向传播
                y_pred = self.model(X)
                
                # 计算损失
                loss = self.loss_fn(y_pred, y)

                # 清空梯度
                self.optimizer.zero_grad()

                # 反向传播
                loss.backward()

                # 参数更新
                self.optimizer.step()

            # 计算当前轮次的训练集和测试集误差
            train_loss = self.get_loss(self.train_dataloader)
            test_loss = self.get_loss(self.test_dataloader)

            print(f"Epoch {epoch}/{epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}")

    def predict(self, X):
        """
        使用训练好的模型进行预测

        参数:
        X (torch.Tensor): 输入数据

        返回:
        torch.Tensor: 预测结果
        """
        self.model.eval()  # 模型进入评估模式
        X = X.to(self.device)  # 将输入数据移动到同一设备上
        with torch.no_grad():  # 禁用梯度计算，因为只做预测
            y_pred = self.model(X)  # 前向传播
        return y_pred

# class Loss(nn.Module):
#     #训练过程监控函数
#     def __init__(self, train_dataloader, model, loss_fn, device):
#         self.train_dataloader = train_dataloader
#         self.model = model
#         self.loss_fn = loss_fn
#         self.device = device
#     def get_loss(self,dataloader,model,loss_fn):
#         with torch.no_grad():
#             #model.eval() #评估模式。
#             losses = []
#             for X,y in dataloader:
#                 #此处可以将数据指定设备。配合前面lr.to（device=device)使用。
#                 X = X.to(device=device)
#                 y = y.to(device=device)
#                 y_pred = model(X) #批量计算。"model的参数形状和X要能对上口型。才能参与矩阵运算。" 。
#                 loss = loss_fn(y_pred,y) #结果也在对应的device上。
#                 losses.append(loss.item())
            # return np.array(losses).mean()

        # #利用梯度下降开始训练
        # def train(model=lr,dataloader=train_dataloader,optimizer=optimizer,loss_fn=loss_fn,epochs=20):
        #     for epoch in range(1,epochs+1):
        #         #每一轮进行训练
        #         for X,y in train_dataloader:
        #             X = X.to(device=device)
        #             y = y.to(device=device)
        #             #正向传播：
        #             y_pred = ________空8__________#这个一调用实际上经过了forward过程。对于多层网络的情况，这样关系是函数嵌套：f3(f2(f1（x）))
        #             #计算损失：
        #             loss =_________空9_________           

        #     #清空梯度：
        #             optimizer.zero_grad()


        #             #反向传播求导：
        #             loss.________空10__________() #loss = f(w,b) 。loss是单值函数。这里就是一个单值函数对一堆变量求偏导。
        #             #优化 step的意思就是减去一次梯度：
        #             optimizer.________空11__________()
        #         #以上操作更新了模型。而且以上每走一步就更新了一次。
        #         print(f"当前是第{epoch}轮，训练集的误差为：{get_loss(dataloader=train_dataloader)},测试集的误差为：{get_loss(dataloader=test_dataloader)}")

if __name__ == '__main__':
    device="cuda:0" if torch.cuda.is_available() else "cpu" 
    df = pd.read_excel('class/the boston_house_prices.xlsx')  
    X = df.drop('MEDV', axis=1).values
    # print(type(X))
    y = df['MEDV'].values

    #切分数据
    X_train,X_test,y_train,y_test = train_test_split(X,y,test_size=0.2,random_state=1)
    y_train = y_train.reshape((-1,1))
    y_test = y_test.reshape((-1,1))

    #进行数据scalar
    mean_ = X_train.mean(axis=0)
    std_ = X_train.std(axis=0)
    #print(mean_.shape)
    #print(X_test.shape)
    X_train= (X_train - mean_) / std_
    X_test = (X_test - mean_) / mean_

    #以下：定义训练集的数据加载器
    train_dataset = MyDataset(X=X_train,y=y_train)

    #测试代码
    #print(train_dataset[0]) #根据idex输出对应的训练样本特征
    #封装为dataloader。Shuffle是指打乱顺序。droplast是把不足的那一批给干掉了。
    train_dataloader = DataLoader(dataset=train_dataset,batch_size=32,shuffle=True,drop_last=True)
    #for X,y in train_dataloader:
    #   print(X.shape,y.shape)

    #以下定义测试集的数据加载器
    test_dataset = MyDataset(X=X_test,y=y_test)
    test_dataloader = DataLoader(dataset=test_dataset,batch_size=32,shuffle=True,drop_last=True)
    lr = LinearRegression()
    #此处可以统一将模型的参数转移到指定的设备上：
    lr.to(device=device)

    #定义损失函数
    loss_fn = nn.MSELoss()
    
    #定义优化器（以下是指把参数lr.parameters托管给优化器SGD）
    optimizer = torch.optim.SGD(params=lr.parameters(),lr=1e-2)

    # 创建Trainer对象并开始训练
    trainer = Trainer(model=lr, train_dataloader=train_dataloader, test_dataloader=test_dataloader, loss_fn=loss_fn, optimizer=optimizer, device='cpu')
    trainer.train(epochs=20)  # 训练20个epoch

    X_new = torch.randn(5, 13)  # 假设有5个样本，每个样本有10个特征
    y_pred = trainer.predict(X_new)
    print(y_pred)  # 打印预测结果

#测试代码
#print(lr.state_dict())
#print(len(lr.state_dict()["linear.weight"][0]))

